package m3f.factorization;

import org.jblas.DoubleMatrix;

import m3f.computing.ConcurrentTemplate;
import m3f.data.MultimodalDataProvider;
import m3f.matrix.Minibatch;

public class ConcurrentOnlineLearner extends ConcurrentTemplate {
	
	private MultimodalDataProvider trainingData;
	private double gamma;
	private double alpha;
	private double beta;
	private double omega;
	private DoubleMatrix lambdaI;
	private DoubleMatrix[] P;
	private DoubleMatrix[] Q;
	
	public ConcurrentOnlineLearner(MultimodalDataProvider data, int factors, double gamma, double lambda, double alpha, double beta, double omega){
		super(data.matrixRows()/data.getMinibatchSize()+1);
		this.trainingData = data;
		this.gamma = gamma;
		this.alpha = alpha;
		this.beta = beta;
		this.omega = omega;
		lambdaI = DoubleMatrix.eye(factors).muli(lambda);
		P = new DoubleMatrix[workers()];
		Q = new DoubleMatrix[workers()];
		for(int i = 0; i < workers(); i++){
			P[i] = DoubleMatrix.rand(data.visualFeatures(), factors).muli(0.01);
			Q[i] = DoubleMatrix.rand(data.textFeatures(), factors).muli(0.01);
		}
		System.out.println("Problem size: " + getProblemSize());
		System.out.println("Workers: " + workers());
	}
	
	public DoubleMatrix getFinalP(){
		DoubleMatrix Pf = DoubleMatrix.zeros(P[0].rows, P[0].columns);
		Double w = new Double(1/(double)workers());
		for(int i = 0; i < w; i++){
			Pf.addi(P[i]);
		}
		Pf.muli(w);
		return Pf;
	}
	
	public DoubleMatrix getFinalQ(){
		DoubleMatrix Qf = DoubleMatrix.zeros(Q[0].rows, Q[0].columns);
		Double w = new Double(1/(double)workers());
		for(int i = 0; i < w; i++){
			Qf.addi(Q[i]);
		}
		Qf.muli(w);
		return Qf;
	}
	
	@Override
	public void singleOperation(int reg) {
		int idx = computeIndex(reg);
		long start = System.currentTimeMillis();
		// Prepare minibatch
		Minibatch vv, tt;
		DoubleMatrix hh, A, B, C;
		synchronized (trainingData) {
			vv = trainingData.nextVisual();
			tt = trainingData.nextTextual();
		}
		// 1. Compute the latent representation for the minibatch:
		hh = latentRepresentation(vv, tt, idx);
		// 2. Update model matrices:
		double omga = 1 - gamma*alpha;
		double omgb = 1 - gamma*beta;
		C = hh.mmul(hh.transpose());
		// Update P
		A = vv.matrix().mmul(hh.transpose()).mul(gamma);
		B = P[idx].mmul(C).mul(-gamma);
		P[idx].muli(omga).addi(A).addi(B);
		// Update Q
		A = tt.matrix().mmul(hh.transpose()).mul(gamma);
		B = Q[idx].mmul(C).mul(-gamma);
		Q[idx].muli(omgb).addi(A).addi(B);
		System.out.println("Minibatch done @ worker "+idx+" (" + (System.currentTimeMillis()-start) + "ms)");
	}
	
	private DoubleMatrix latentRepresentation(Minibatch V, Minibatch T, int idx){
		DoubleMatrix A, G;
		G = P[idx].transpose().mmul(P[idx]);
		A = Q[idx].transpose().mmul(Q[idx]);
		G.addi(A).addi(lambdaI);
		return OnlineJointFactorization.latentRepresentation(V, T, P[idx], Q[idx], G, omega);
	}
	
	private int computeIndex(int reg){
		double idx = ((double)reg/(double)getProblemSize())*(double)workers();		
		return (int)Math.floor(idx);
	}
}
